# see also: http://mail.python.org/pipermail/python-dev/2008-January/076194.html

def redefine_method(original_method, new_method):

    all_modules = copy(sys.modules)
    for module_name, module in all_modules.iteritems():
        # FIXME: this "magically" makes it work (doing load_module), but it is slow...
        try:
            load_module(module_name, *find_module(module_name))
        except ImportError:
            # TODO: couldn't load the module, keep going?
            continue
        for obj_name in dir(module):
            obj = getattr(module, obj_name)
            if hasattr(obj, "func_code") and obj != new_method:
                if obj.func_code == original_method.func_code:
                    setattr(module, obj_name, new_method)


class ScopedMock(object):
    
    def __init__(self, real_object, mock_object):
        
        if hasattr(real_object, "im_func") and hasattr(mock_object, "im_func"):
            
            # mock instance methods
            self.__original_func_code = real_object.im_func.func_code
            def mock():
                real_object.im_func.func_code = mock_object.im_func.func_code
            def unmock():
                real_object.im_func.func_code = self.__original_func_code
            self.__mock = mock
            self.__unmock = unmock
        
        elif hasattr(real_object, "func_code") and hasattr(mock_object, "func_code"):
            
            # mock functions
            self.__original_func_code = real_object.func_code
            def mock():
                real_object.func_code = mock_object.func_code
            def unmock():
                real_object.func_code = self.__original_func_code
            self.__mock = mock
            self.__unmock = unmock
        
        elif hasattr(real_object, "__dict__") and hasattr(mock_object, "__dict__"):
            
            # mock entire objects
            self.__mock_children = []
            for attr_name in real_object.__dict__:
                mock_child = ScopedMock(
                    getattr(real_object, attr_name),
                    getattr(mock_object, attr_name),
                )
                self.__mock_children.append( mock_child )
            def mock():
                for mock_child in self.__mock_children:
                    mock_child.mock()
            def unmock():
                for mock_child in self.__mock_children:
                    mock_child.unmock()
            self.__mock = mock
            self.__unmock = unmock
            
        else:
            
            # don't try to mock other things
            self.__mock = lambda:None
            self.__unmock = lambda:None
    
    def mock(self):
        self.__mock()
    
    def unmock(self):
        self.__unmock()
    
    def __enter__(self):
        self.__mock()
    
    def __exit__(self, *args):
        self.__unmock()
